import tiktoken

from .lang import detect_language, language_name
from .api import chat_completion_create

def cut_content(content: str, length: int, model: str = 'gpt-3.5-turbo', both_side: bool = False):
    if length <= 0:
        return content, None
    enc = tiktoken.encoding_for_model(model)
    tokens = enc.encode_ordinary(content)
    if len(tokens) <= length:
        return content, None
    truncation = (len(tokens), length)
    if both_side:
        omit_token = enc.encode_ordinary('...')
        tokens = tokens[:length//2] + omit_token + tokens[-length//2:]
    else:
        tokens = tokens[:length]
    content = ''
    accum = b''
    for part in enc.decode_tokens_bytes(tokens):
        accum += part
        if len(accum) >= 8:
            accum = b''
        try:
            decoded = accum.decode()
        except UnicodeDecodeError:
            pass
        else:
            content += decoded
            accum = b''
    return content, truncation

def omit_content_line(content: str, n: int, size_mark: bool = False, both_side: bool = True) -> str:
    content = content.strip().replace('\n', ' ')
    while content.count('  '):
        content = content.replace('  ', ' ')
    content, truncation = cut_content(content, n, both_side=both_side)
    if truncation:
        if not both_side:
            content = content + '...'
        if size_mark:
            content = content + f' ({truncation[0]} tokens)'
    return content

def list_conversation(messages: list[dict], include_system=True, include_tools=True, list_number=True, size_mark=False, limit_tokens=32, slice_history=0) -> str:
    history = []
    history_system = []
    for i, m in enumerate(messages):
        prefix = str(i) + '. ' if list_number else ''
        if m['role'] == 'system' and include_system:
            history_system.append(prefix + 'System: ' + omit_content_line(m['content'], 60, size_mark))
        elif m['role'] == 'user':
            history.append(prefix + 'User: ' + omit_content_line(m['content'], limit_tokens, size_mark))
        elif m['role'] == 'assistant':
            if m.get("content") and m.get("tool_calls") and include_tools:
                history.append(prefix + 'Tool call: ' + '; '.join(x['function']['name'] + ' ' + omit_content_line(x['function']['arguments'], limit_tokens, True) for x in m['tool_calls']) + '; Reason: ' + omit_content_line(m['content'], limit_tokens, size_mark))
            elif m.get("tool_calls") and include_tools:
                history.append(prefix + 'Tool call: ' + '; '.join(x['function']['name'] + ' ' + omit_content_line(x['function']['arguments'], limit_tokens, True) for x in m['tool_calls']))
            elif m.get("content"):
                history.append(prefix + 'Assistant: ' + omit_content_line(m['content'], limit_tokens, size_mark))
        elif m['role'] == 'tool' and include_tools:
            history.append(prefix + 'Tool call result: ' + omit_content_line(m['content'], limit_tokens, size_mark))
    if slice_history > 0:
        history = history[-slice_history:]
    history = '\n'.join(history_system + history)
    return history

def conversation_choices(messages: list[dict], num_choices: int | str = '1~3') -> list[str]:
    history = list_conversation(messages, include_tools=False, list_number=False, limit_tokens=80, slice_history=3)
    language = language_name(detect_language(history))
    prompt = '''Based on the following conversation:\n{history}\nGenerate {num_choices} possible user replies separated by newlines. Output in {language}. Output the replies in quotes ("").'''
    prompt = prompt.format(
        history=history,
        language=language,
        num_choices=num_choices,
    )
    completion = chat_completion_create(
        messages=[{'role': 'user', 'content': prompt}],
        model='gpt-3.5-turbo-0125',
        temperature=1,
        max_tokens=128,
    )
    result = list()
    for choice in completion.choices:
        response = choice.message.content
        if response:
            # print(prompt)
            # print(response)
            for i, res in enumerate(response.splitlines()):
                res = res.strip('"')
                if res.startswith(f'{i+1}. '):
                    res = res[res.index('.') + 1:]
                res = res.strip().strip('"')
                if res:
                    result.append(res)
    return result

def conversation_title(messages: list[dict]):
    history = list_conversation(messages, include_tools=False, include_system=False, list_number=False, limit_tokens=80, slice_history=4)
    language = language_name(detect_language(history))
    prompt = '''Generate a unique and memorizable {language} title for the following conversation:\n{history}\nOutput the title in quotes ("").'''
    prompt = prompt.format(
        history=history,
        language=language,
    )
    completion = chat_completion_create(
        messages=[{'role': 'user', 'content': prompt}],
        model='gpt-3.5-turbo-0125',
        temperature=1,
        max_tokens=64,
    )
    response = completion.choices[0].message.content
    if response:
        return response.strip().strip('"')
    else:
        return 'Untitled'
